#!/usr/bin/python
# 
# Copyright (C) 2009 Brian M Hunt
#
# Released under the "MIT" license, as follows:
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the "Software"), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
# 
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.
#
#
# Requires Python 2.6 on Mac OS X because of 
# Problem: http://python.org/sf/1092502
#
#
# See RFC 1521 re. message/partial re-encodings
# http://webee.technion.ac.il/labs/comnet/netcourse/CIE/RFC/1521/24.htm
#

import os, errno
import logging
import imaplib
import rfc822
import time
import sys
import re

from datetime import date
from optparse import OptionParser
from subprocess import Popen, PIPE
from email import parser, feedparser

# Filename for logs
_email_parser = parser.Parser()


def save_message(email, dated_ctime, dest, filename):
    """Save an email to ``dest``
    email is presumed to be application/pdf or other saveable file
    """
    if not filename:
        logging.error("Message part does not have a filename")
        return

    full_path = os.path.join(dest, filename)

    encoding = email['Content-Transfer-Encoding']

    ### SECURITY: CHECK filename doesn't contain ".."/etc
    logging.info("Saved %s." % full_path)

    try:
        f = open(full_path, 'w')
        f.write(email.get_payload(decode=True))
        f.close()
        # set appropriate modification / creation times
        os.utime(full_path, (dated_ctime, dated_ctime))
    except Exception, e:
        logging.error("Could not write file: %s." % e)
        raise

def _part_filename(partid, partno):
    return "%s.%d" % (partid, partno)


def save_partial(eml, dest, uudeview):
    """Save the parts of an email. Recompile them if all the parts are together.
    """
    content_string = eml.get('Content-Type')
    match = re.search('id="(.*?)";.*number=(\d+);.*total=(\d+)', content_string, re.M|re.S)

    if not match:
        raise Exception("Invalid partial content string: %s", content_string)

    part_id     = match.group(1)
    part_number = int(match.group(2))
    part_total  = int(match.group(3))

    filename = _part_filename(part_id, part_number)
    full_path = os.path.join(dest, filename)

    logging.debug("%s [%d of %d]." % (full_path, part_number, part_total))
    
    written = False
    with open(full_path, "w") as f:
        f.write(eml.as_string())

    # for message 1 of N, there's a content-header e.g.
    # Content-Type: application/pdf
    # for messages 2 to N, there is only text/plain.
    # However, message 1 also contains text/plain != the application content
#    if part_number == 1:
#        for part in eml.walk():
#            if part.get_content_type() == 'application/pdf':
#                with open(filename, "w") as f:
#                    f.write(part.as_string(True))
#                    written = True
#    else:
#        for part in eml.walk():
#            if part.get_content_type() == 'text/plain':
#                with open(filename, "w") as f:
#                    f.write(part.as_string(True))
#                    written = True


    part_files = [os.path.join(dest, _part_filename(part_id,pno))
        for pno in xrange(1, part_total + 1)]

    logging.debug("Part files: %s " % ",".join(part_files))

    if all(os.path.exists(partfile) for partfile in part_files):
        logging.debug("Got all parts - reconstructing %s." % part_id)

        cmd = [uudeview, '-i', '-n'] + part_files

        logging.debug("Calling command: %s" % (" ".join(cmd)))

        p = Popen(cmd, stdout=PIPE, stdin=PIPE, stderr=PIPE)
            # , stdin=PIPE, stdout=PIPE, stderr=PIPE)

        p.communicate()
        # this creates 0001.txt
        # p.communicate("\n".join(part_files))

        #parser = feedparser.FeedParser()

        # grep out the stupid --DC_BOUND_PRE in the middle
        with open("0001.txt") as fin:
            with open("0001.msg", 'w') as fout:
                i = 0;

                # skip empty lines and lines that start with --DC_BOUND_PRE
                for line in fin:
                    i = i + 1
                    if i > 1000:
                        if not line or line.find("--DC_BOUND_PRE") == 0:
                            last_line = line
                            continue
                    fout.write(line)

                fout.write(last_line)

        cmd = [uudeview, '-i', '-n', '0001.msg']

        p = Popen(cmd, stdout=PIPE, stdin=PIPE, stderr=PIPE)
            # , stdin=PIPE, stdout=PIPE, stderr=PIPE)

        p.communicate()

        for fname in part_files + ["0001.txt", "0001.msg"]:
            logging.debug("Removing %s" % fname)
            os.unlink(fname)

        logging.info("Compiled multi-part message")
    else:
        logging.info("Multi-part message is missing parts.")
        
        #filename = message.get_filename()
        #dated_ctime = time.mktime(rfc822.parsedate(message['Date']))

        #logging.debug("Writing file %s" % filename)

        #for part in message.walk():
        #     part_type = part.get_content_type()
        #    if part_type == 'application/pdf':
        #        save_message(part, dated_ctime, dest, filename)

        #full_path = os.path.join(dest, message.get_filename())
        #with open(full_path, "w") as f:
        #    f.write(message.get_payload(decode=True))
#    
    ## Now check for all the parts, and if all parts are here, recombine them
    # and decode the email
    


def slurp_msg(M, num, dest, uudeview):
    """
    Take an imap object and message number and read it and write the attachments to disk
    """
    logging.info("Retrieving message #%s." % num)

    status, msg = M.fetch(num, '(RFC822)')
    if status != 'OK':
        logging.error("Fetch [msg #%s] status: " % (num, status))
        raise Exception("Improper select status")

    email = _email_parser.parsestr(msg[0][1])

    #if email['Subject'].find('part') == -1:
    #    logging.debug("Skipping non-multipart message")
    #    return

    logging.debug("Email dated %s from %s, subject: %s" % (
                email['Date'],
                email['From'],
                email['Subject'])
            )

    dated_ctime = time.mktime(rfc822.parsedate(email['Date']))

    if email.get_content_type() == 'message/partial':
        save_partial(email, dest, uudeview)

    else:
        for part in email.walk():
            part_type = part.get_content_type()

            if part_type == 'application/pdf':
                filename = part.get_filename()
                logging.debug("Email part Filename: %s" % filename)
                save_message(part, dated_ctime, dest, filename)

    logging.debug("Marking message %s as seen" % num)

    status, flags = M.store(num,"+FLAGS","(\\Seen)")
    if status != 'OK':
        # non-fatal
        logging.error("Could not mark msg #%s seen: " % (num, status))

    # rebuild parts, if possible


def connection(host, port, user, password):
    logging.debug("Connecting to %s:%d." % (host, port))
    M = imaplib.IMAP4_SSL(host, port)

    try:
        logging.debug("Logging in as %s" % user)
        M.login(user, password)
    except Exception, e:
        logging.exception("Unable to log in: %s" % e)
        raise

    return M



def main():
    parser = OptionParser()

    parser.add_option("-o", "--output-dir", action="store", type="string",
        dest="output_dir", default="", help="Directory where attachments are saved")

    parser.add_option("-s", "--server", action="store", type="string",
        dest="host", default="localhost", help="Host with the IMAP server")

    parser.add_option("-p", "--port", action="store", type="int",
        dest="port", default="993", help="Port of the IMAP server")

    parser.add_option("-u", "--user", action="store", type="string",
        dest="user", default="root", help="User for the IMAP account")

    parser.add_option("-z", "--password", action="store", type="string",
        dest="password", default="", help="User's IMAP password")

    parser.add_option("-v", "--verbose", action="store_true",
        dest="verbose", default="", help="Verbose output")

    parser.add_option("-l", "--log", action="store",
        dest="log_file", default="", help="Log file")

    parser.add_option("-D", "--dated-dir", action="store_true",
        dest="dated_dir", default="", help="Create a directory of year-month for the files")

    parser.add_option("--uudeview", action="store",
        dest="uudeview", default="/usr/local/bin/uudeview", help="path to uudeview command")

    parser.add_option("-b", "--box", action="store", type="string",
        dest="mailbox", default="INBOX",
        help="Mailbox for getting the messages (default 'INBOX')")

    parser.add_option("-m", "--messages", action="store", type="string",
        dest="messages", default="",
        help="Comma separated list of message numbers to download (default: get all unseen messages)")

    (options, args) = parser.parse_args()

    if options.verbose:
        log_level = logging.DEBUG
    else:
        log_level = logging.INFO

    if not options.output_dir:
        options.output_dir = '.'

    if options.dated_dir:
        today = date.today()
        output_dir = os.path.join(options.output_dir,
            "%d-%02d" % (today.year, today.month))
    else:
        output_dir = options.output_dir

    try:
        if output_dir:
            os.makedirs(output_dir)
    except OSError, exc:
        if exc.errno == errno.EEXIST:
            pass
        else: raise

    os.chdir(output_dir)

    if options.log_file:
        logging.basicConfig(level=log_level,
            filename=os.path.join(output_dir, options.log_file),
            format="%(asctime)s - %(levelname)s - %(message)s",)
    else:
        logging.basicConfig(level=log_level,
            format="%(asctime)s - %(levelname)s - %(message)s",)
    
    M = connection(options.host, options.port,
        options.user, options.password)

    status, count = M.select(options.mailbox)

    if status != 'OK':
        logging.error("Select status: " % status)
        raise Exception("Improper select status")

    if options.messages:
        message_list = options.messages.split(",")
    else:
        # message list is set of all unseen messages
        message_list = M.search(None, 'UNSEEN')[1][0].split()
        if message_list:
            logging.info("%s new messages." % len(message_list))


    for num in message_list:
        try:
            slurp_msg(M, num, output_dir, options.uudeview)
        except Exception, e:
            logging.exception("Unable to slurp message %s: %s", num, e)
            raise

    M.close()

if __name__ == '__main__':
    main()

